#include "sample_svp_nnie_software.h"

#include <time.h>

# include "efficientdet.h"

// 参数设置
static SAMPLE_SVP_NNIE_MODEL_S stEfficientModel = {0}; // 模型参数
static SAMPLE_SVP_NNIE_PARAM_S stEfficientNnieParam = {0}; // NNIE参数
static SVP_SOFT_EfficientDet stEfficientSoftParam = {0}; // soft infor

static HI_U8 *JPGfile_pu8Buf = NULL;  // 存储jpg文件的malloc空间
static HI_U32 u32VdecChnNum = 1; // vdec通道数
static VPSS_GRP VpssGrp; // vpss 组数 用于断开vpss
static HI_BOOL abChnEnable[VPSS_MAX_CHN_NUM]; // VPSS 通道使能 标志位
static HI_S32 s32VpssGrp = 0; // vpss group号 设置 
static HI_S32 as32VpssChn[] = {VPSS_CHN0, VPSS_CHN1}; // vpss 通道号 设置
static VIDEO_FRAME_INFO_S stDetectInfo; // 检测图像


/* *****************************************************

jpg -> txt 退出函数

******************************************************* */
static void EfficientDet_Exit(int end_select){
    HI_S32 ii;
    HI_S32 s32Ret = HI_SUCCESS;

    if(JPGfile_pu8Buf != NULL){  // 清理存储jpg文件的malloc内存
        free(JPGfile_pu8Buf);
        JPGfile_pu8Buf = NULL;
    }

    switch (end_select){
        case 1:
            goto END1; // 退出sys+vb初始化
        case 2:
            goto END2; // 退出vdec设置的vb
        case 3:
            goto END3; // 停止vdec
        case 4:
            goto END4; // 关闭vpss
        case 5:
            goto END5; // 断开 vpss vdec
        case 6:
            goto Model_END6; // 模型参数内存释放
        case 7:
            goto Model_END7; // NNIE硬件释放
        case 8:
            goto Model_END8; // 释放soft
        case 9:
            goto END9; // 释放VPSS通道的数据
        
        default:
            break;
    }

END9: // 释放VPSS通道的数据
    s32Ret = HI_MPI_VPSS_ReleaseChnFrame(s32VpssGrp, as32VpssChn[0], &stDetectInfo);
    if (HI_SUCCESS != s32Ret)
        printf("Error(%#x),HI_MPI_VPSS_ReleaseChnFrame failed,Grp(%d) chn(%d)!\n", s32Ret, s32VpssGrp, as32VpssChn[0]);

Model_END8: // 释放soft
    Efficient_Soft_Deinit(&stEfficientSoftParam);

Model_END7: // NNIE硬件释放
    s32Ret = SAMPLE_COMM_SVP_NNIE_ParamDeinit(&stEfficientNnieParam);
    SAMPLE_SVP_CHECK_EXPR_TRACE(HI_SUCCESS != s32Ret,SAMPLE_SVP_ERR_LEVEL_ERROR, "Error,SAMPLE_COMM_SVP_NNIE_ParamDeinit failed!\n");
    memset(&stEfficientNnieParam, 0, sizeof(SAMPLE_SVP_NNIE_PARAM_S));

Model_END6: // 模型参数内存释放
    s32Ret = SAMPLE_COMM_SVP_NNIE_UnloadModel(&stEfficientModel);
    SAMPLE_SVP_CHECK_EXPR_TRACE(HI_SUCCESS != s32Ret,SAMPLE_SVP_ERR_LEVEL_ERROR,
        "Error,SAMPLE_COMM_SVP_NNIE_UnloadModel failed!\n");
    memset(&stEfficientModel, 0, sizeof(SAMPLE_SVP_NNIE_MODEL_S));

END5: // 断开 vpss vdec
    for(ii=0; ii<u32VdecChnNum; ii++){
        s32Ret = SAMPLE_COMM_VDEC_UnBind_VPSS(ii, ii); // 使vdec和vpss断连
        if(s32Ret != HI_SUCCESS) SAMPLE_PRT("vdec unbind vpss fail for %#x!\n", s32Ret);
    }

END4: // 关闭vpss
    for(ii=VpssGrp; ii>=0; ii--){ // i不能为unsigned int型
        VpssGrp = ii;
        SAMPLE_COMM_VPSS_Stop(VpssGrp, &abChnEnable[0]); // 关闭所有group的vpss
    }

END3: // 停止vdec
    SAMPLE_COMM_VDEC_Stop(u32VdecChnNum);

END2: // 退出vdec设置的vb
    SAMPLE_COMM_VDEC_ExitVBPool();

END1: // 退出sys+vb初始化
    SAMPLE_COMM_SYS_Exit();

exit(0);
}

// 释放内存
void EfficientDet_JPG2txt_HandleSig(){
    EfficientDet_Exit(0);
}

/* *****************************************************

jpg -> txt 主函数

******************************************************* */

void EfficientDet_JPG2txt()
{
    // ls imgs | sed "s:^:./efficient/data/imgs/:" >> img_list.txt
    FILE *fp_list = fopen("./efficient/data/img_list.txt", "r");// 判断list是否为空
    if(fp_list == NULL){ // list不能为空
        printf("Error: no file\n"); 
        exit(1);
    } 
    HI_CHAR *pcModelName = "./efficient/data/efficient_red_car.wk";
    start_reckon_by_time(); // 开始计时

    // ***************       1 sys+vb init       ***************
    HI_S32 ii = 0;
    HI_S32 s32Ret = HI_SUCCESS;
    {
        VB_CONFIG_S stVbConfig;
        printf("step1: vb init...\n");
        memset(&stVbConfig, 0, sizeof(VB_CONFIG_S));
        stVbConfig.u32MaxPoolCnt             = 1; // 系统可容纳缓存池个数
        stVbConfig.astCommPool[0].u32BlkCnt  = 10*u32VdecChnNum; // 每个缓存池的缓存块个数
        stVbConfig.astCommPool[0].u64BlkSize = COMMON_GetPicBufferSize(1920, 1080,
                                               PIXEL_FORMAT_YVU_SEMIPLANAR_420, DATA_BITWIDTH_8, COMPRESS_MODE_SEG, 0); // 缓存块大小
                                                    
        s32Ret = SAMPLE_COMM_SYS_Init(&stVbConfig);
        if(s32Ret != HI_SUCCESS){
            SAMPLE_PRT("init sys fail for %#x!\n", s32Ret);
            EfficientDet_Exit(1);
        }
    }
    PCPrintTime("Sys init", 0, 1, 1);
    
    /********    2 vdec的vb init   *********/
    SAMPLE_VDEC_ATTR astSampleVdec[VDEC_MAX_CHN_NUM];
    {
        printf("step2: vdec vb init...\n");
        for(ii=0; ii<u32VdecChnNum; ii++){ // 有一个vdec通道
            astSampleVdec[ii].enType                            = PT_JPEG;
            astSampleVdec[ii].u32Width                          = 3840;
            astSampleVdec[ii].u32Height                         = 2160;
            astSampleVdec[ii].enMode                            = VIDEO_MODE_FRAME;
            astSampleVdec[ii].stSapmleVdecPicture.enPixelFormat = PIXEL_FORMAT_YVU_SEMIPLANAR_420;
            astSampleVdec[ii].stSapmleVdecPicture.u32Alpha      = 255;
            astSampleVdec[ii].u32DisplayFrameNum                = 2;
            astSampleVdec[ii].u32FrameBufCnt                    = astSampleVdec[ii].u32DisplayFrameNum + 1;
        }
        s32Ret = SAMPLE_COMM_VDEC_InitVBPool(u32VdecChnNum, &astSampleVdec[0]); // 根据类型、大小,设置公共视频缓存池属性(3个API)
        if(s32Ret != HI_SUCCESS){
            SAMPLE_PRT("init mod common vb fail for %#x!\n", s32Ret);
            EfficientDet_Exit(2);
        }
    }
    PCPrintTime("vdec vb init", 0, 1, 1);

    /********    3 开启vdec    *********/
    {
        printf("step3: start vdec...\n");
        s32Ret = SAMPLE_COMM_VDEC_Start(u32VdecChnNum, &astSampleVdec[0]);
        if(s32Ret != HI_SUCCESS){
            SAMPLE_PRT("start VDEC fail for %#x!\n", s32Ret);
            EfficientDet_Exit(3);
        }
    }
    PCPrintTime("start vdec", 0, 1, 1);

    /********    4 开启vpss    *********/
    {
        printf("step4: start vpss...\n");
        // 用于创建vpss的属性
        VPSS_GRP_ATTR_S stVpssGrpAttr; // vpss组属性
        stVpssGrpAttr.u32MaxW = 3840; // 3840 输入图像
        stVpssGrpAttr.u32MaxH = 2160; // 2160
        stVpssGrpAttr.stFrameRate.s32SrcFrameRate = -1; // 组帧率
        stVpssGrpAttr.stFrameRate.s32DstFrameRate = -1;
        stVpssGrpAttr.enDynamicRange = DYNAMIC_RANGE_SDR8; // 输入图像动态范围
        stVpssGrpAttr.enPixelFormat  = PIXEL_FORMAT_YVU_SEMIPLANAR_420; // 输入图像像素格式
        stVpssGrpAttr.bNrEn   = HI_FALSE; // NR 使能开关
        // 给VPSS哪些通道使能 标志位
        memset(abChnEnable, 0, sizeof(abChnEnable));

        VPSS_CHN_ATTR_S astVpssChnAttr[VPSS_MAX_CHN_NUM]; // vpss通道属性
        astVpssChnAttr[0].u32Width                    = EFFICIENTDET_INPUT_WIDTH; // 输出图片宽/高
        astVpssChnAttr[0].u32Height                   = EFFICIENTDET_INPUT_HEIGHT;
        for(ii = 0; ii < 1; ii++){
            abChnEnable[ii] = HI_TRUE; // 第0个vpss通道使能控制位

            astVpssChnAttr[ii].enChnMode                   = VPSS_CHN_MODE_USER; // VPSS_CHN_MODE_AUTO
            astVpssChnAttr[ii].enCompressMode              = COMPRESS_MODE_NONE; // 目标图像压缩模式 COMPRESS_MODE_SEG
            astVpssChnAttr[ii].enDynamicRange              = DYNAMIC_RANGE_SDR8; // 目标图像动态范围
            astVpssChnAttr[ii].enPixelFormat               = PIXEL_FORMAT_YVU_SEMIPLANAR_420; // 目标图像像素格式
            astVpssChnAttr[ii].stFrameRate.s32SrcFrameRate = -1; // 帧率控制信息
            astVpssChnAttr[ii].stFrameRate.s32DstFrameRate = -1;
            astVpssChnAttr[ii].u32Depth                    = 1; // 用户获取通道图像的队列长度 0
            astVpssChnAttr[ii].bMirror                     = HI_FALSE; // 水平镜像使能
            astVpssChnAttr[ii].bFlip                       = HI_FALSE; // 垂直翻转使能
            astVpssChnAttr[ii].stAspectRatio.enMode        = ASPECT_RATIO_NONE; // 幅形比参数
            astVpssChnAttr[ii].enVideoFormat               = VIDEO_FORMAT_LINEAR; // 目标图像视频格式
        }

        for(ii=0; ii<1; ii++){ // 1个通道 
            VpssGrp = ii;
            s32Ret = SAMPLE_COMM_VPSS_Start(VpssGrp, &abChnEnable[0], &stVpssGrpAttr, &astVpssChnAttr[0]); // 开启vpss通道
            if(s32Ret != HI_SUCCESS){
                SAMPLE_PRT("start VPSS fail for %#x!\n", s32Ret);
                EfficientDet_Exit(4);
            }
        }      
    }
    PCPrintTime("start vpss", 0, 1, 1);

    /********    5 vdec连接vpss    *********/
    {
        printf("step5: vdec link vpss...\n");
        for(ii=0; ii<u32VdecChnNum; ii++){
            s32Ret = SAMPLE_COMM_VDEC_Bind_VPSS(ii, ii);
            if(s32Ret != HI_SUCCESS){
                SAMPLE_PRT("vdec bind vpss fail for %#x!\n", s32Ret);
                EfficientDet_Exit(5);
            }
        }
        usleep(10);  
    }
    PCPrintTime("link vdec to vpss", 0, 1, 1);
          
    /********    6 发送数据到vdec    *********/
    VDEC_THREAD_PARAM_S stVdecSend[VDEC_MAX_CHN_NUM]; // vdec参数
    {
        printf("step6: set vdec para...\n");
        for(ii=0; ii<u32VdecChnNum; ii++){ 
            snprintf(stVdecSend[ii].cFileName, sizeof(stVdecSend[ii].cFileName), "1080P.jpg");
            snprintf(stVdecSend[ii].cFilePath, sizeof(stVdecSend[ii].cFilePath), ".");
            stVdecSend[ii].enType          = astSampleVdec[ii].enType;
            stVdecSend[ii].s32StreamMode   = astSampleVdec[ii].enMode;
            stVdecSend[ii].s32ChnId        = ii;
            stVdecSend[ii].s32IntervalTime = 1000;
            stVdecSend[ii].u64PtsInit      = 0;
            stVdecSend[ii].u64PtsIncrease  = 0;
            stVdecSend[ii].eThreadCtrl     = THREAD_CTRL_START;
            stVdecSend[ii].bCircleSend     = HI_TRUE;
            stVdecSend[ii].s32MilliSec     = 0;
            stVdecSend[ii].s32MinBufSize   = (astSampleVdec[ii].u32Width * astSampleVdec[ii].u32Height * 3)>>1;
        } 
    }
    PCPrintTime("send data to vdec", 0, 1, 1);  

    /************************************************
    step: SVP
    *************************************************/

    // ***************       7 load wk       ***************
    {
        printf("step7: load model...\n");
        s32Ret = SAMPLE_COMM_SVP_NNIE_LoadModel(pcModelName,&stEfficientModel);
        if(s32Ret != HI_SUCCESS){
            SAMPLE_PRT("Error,NNIE_LoadModel failed for %#x!\n", s32Ret);
            EfficientDet_Exit(6);
        } 
    }
    PCPrintTime("load wk", 0, 1, 1);  
    
    // ***************       8 NNIE init       ***************
    {
        printf("step8: NNIE parameter init...\n");
        SAMPLE_SVP_NNIE_CFG_S   stNnieCfg = {0};
        stNnieCfg.u32MaxInputNum = 1; // max input image num in each batch
        stNnieCfg.u32MaxRoiNum = 0;
        stNnieCfg.aenNnieCoreId[0] = SVP_NNIE_ID_0; // set NNIE core
        stEfficientNnieParam.pstModel = &stEfficientModel.stModel;
        s32Ret = SAMPLE_COMM_SVP_NNIE_ParamInit(&stNnieCfg, &stEfficientNnieParam);
        if(s32Ret != HI_SUCCESS){
            SAMPLE_PRT("Error,NNIE_ParamInit failed for %#x!\n", s32Ret);
            EfficientDet_Exit(7);
        } 
    }
    PCPrintTime("NNIE init", 0, 1, 1);  

    // ***************       9 soft init       ***************'
    {
        printf("step9: soft parameter init...\n");
        s32Ret = EfficientDet_Soft_ParamInit(&stEfficientSoftParam);
        if(s32Ret != HI_SUCCESS){
            SAMPLE_PRT("Error,NNIE_ParamInit failed for %#x!\n", s32Ret);
            EfficientDet_Exit(8);
        }
    }
    PCPrintTime("soft init", 1, 1, 1);

    // JPG存储空间
    HI_CHAR cStreamFile[1023]; // 用来存储jpg的名字
    JPGfile_pu8Buf = malloc(stVdecSend[0].s32MinBufSize); // 存储JPG原始数据
    if(JPGfile_pu8Buf == NULL){
        printf("can't alloc %d in send stream thread!\n", stVdecSend[0].s32MinBufSize);
        EfficientDet_Exit(8);
    }
    fflush(stdout); // 标准输出流的清理 打印数据到屏幕上

    int img_count = 0;  // 检测图片数量
    SAMPLE_SVP_NNIE_INPUT_DATA_INDEX_S stInputDataIdx = {0}; // NNIE输入数据位置
    SAMPLE_SVP_NNIE_PROCESS_SEG_INDEX_S stProcSegIdx = {0};  // NNIE处理段id
    start_reckon_by_time(); // 重新计时
    while(1)
    {
        // ***************       向VDEC发送JPG码流       ***************
        VDEC_STREAM_S stStream; // 发送到vdec的数据流
        s32Ret = Efficient_Read_JPG(fp_list, JPGfile_pu8Buf, cStreamFile, &stStream, stVdecSend); // 读取JPG,配置vdec参数
        if(HI_SUCCESS != s32Ret) EfficientDet_Exit(8); // 跳出循环
        img_count += 1;
        printf("\n%dth img name: %s\n", img_count, cStreamFile);

SendAgain: // 向视频解码通道发送码流数据
        s32Ret=HI_MPI_VDEC_SendStream(stVdecSend[0].s32ChnId, &stStream, stVdecSend[0].s32MilliSec); // 0非阻塞
        if(HI_SUCCESS != s32Ret){ // 发送失败
            usleep(stVdecSend[0].s32IntervalTime);
            goto SendAgain;
        }
        PCPrintTime("send stream to vdec", 0, 1, 1);
        
        // 创建txt
        char result_file_name[100];
        sprintf(result_file_name, EFFICIENTDET_SAVE_NAME, img_count);
        stEfficientSoftParam.fp_result = fopen(result_file_name, "w"); // 创建bbox文件   sudo chown ccc bbox_result.txt

        // ***************       从VPSS获取YUV       ***************
        s32Ret = HI_MPI_VPSS_GetChnFrame(s32VpssGrp, as32VpssChn[0], &stDetectInfo, 20); // 获取图像等待时间:20
        if(HI_SUCCESS != s32Ret){
            printf("Error(%#x),HI_MPI_VPSS_GetChnFrame failed, VPSS_GRP(%d), VPSS_CHN(%d)!\n", s32Ret,s32VpssGrp, as32VpssChn[0]);
            goto DETECT_RELEASE;
        }
        PCPrintTime("get yuv from vdec", 0, 1, 1);

        // ***************       yuv420sp -> bgr       ***************
        s32Ret = Efficient_YUV2BGR(&stDetectInfo, &stEfficientNnieParam, &stInputDataIdx);
        if(HI_SUCCESS != s32Ret) goto DETECT_RELEASE;
        PCPrintTime("yuv2bgr", 0, 1, 1);
        PCPrintTime("yuv2bgr", 1, 0, 0); // 只计推理时间

        // ***************       NNIE forward       ***************
        stProcSegIdx.u32SegIdx = 0;
        s32Ret = EfficientDet_NNIE_Forward(&stEfficientNnieParam, &stInputDataIdx, &stProcSegIdx, HI_TRUE);
        if(HI_SUCCESS != s32Ret){
            printf("Error!! NNIE Forward , code:%d\n", s32Ret);
            goto DETECT_RELEASE;
        }
        PCPrintTime("NNIE", 0, 1, 1);

        // ***************       soft forward       ***************
        s32Ret = EfficientDet_Soft(&stEfficientNnieParam, &stEfficientSoftParam);
        if(HI_SUCCESS != s32Ret){
            printf("Error!! soft Forward , code:%d\n", s32Ret);
            goto DETECT_RELEASE;
        }
        PCPrintTime("soft", 0, 1, 1);
        
        // ***************       打印       ***************
        HI_FLOAT* final_result_tmp = (HI_FLOAT*)stEfficientSoftParam.ps32FinalBbox.u64VirAddr;
        fprintf(stEfficientSoftParam.fp_result, "%s\n", cStreamFile);
        for(HI_U32 ii=0; ii<stEfficientSoftParam.final_detection_num; ++ii){ // [score, class, x1, y1, x2, y2]
            if(!EFFICIENTDET_PRINT_FLAG)
                printf("   score: %.4f   class: %d   xyxy: %.4f  %.4f  %.4f  %.4f\n",  
                    final_result_tmp[ii*6], (HI_S32)final_result_tmp[ii*6+1], final_result_tmp[ii*6+2], 
                    final_result_tmp[ii*6+3], final_result_tmp[ii*6+4], final_result_tmp[ii*6+5]);
            fprintf(stEfficientSoftParam.fp_result, "%.4f %d %.1f %.1f %.1f %.1f\n", final_result_tmp[ii*6], (HI_S32)final_result_tmp[ii*6+1], 
                    final_result_tmp[ii*6+2], final_result_tmp[ii*6+3], final_result_tmp[ii*6+4], final_result_tmp[ii*6+5]); 
        } 
        PCPrintTime("print", 1, 1, 1);

    // 释放检测图片
    DETECT_RELEASE:
        s32Ret = HI_MPI_VPSS_ReleaseChnFrame(s32VpssGrp, as32VpssChn[0], &stDetectInfo);
        if (HI_SUCCESS != s32Ret){
            printf("Error(%#x),HI_MPI_VPSS_ReleaseChnFrame failed,Grp(%d) chn(%d)!\n", s32Ret, s32VpssGrp, as32VpssChn[0]);
            EfficientDet_Exit(8);
        }
        Efficient_mean_time(img_count);
    }

    EfficientDet_Exit(8);  // 退出程序
}

